from selenium import webdriver 
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.common.by import By
from selenium.webdriver.edge.options import Options
from selenium.common.exceptions import TimeoutException  # 导入 TimeoutException
from selenium.common.exceptions import ElementClickInterceptedException
from selenium.common.exceptions import StaleElementReferenceException
import pickle
import os
import logging
from datetime import datetime
import argparse
import json

"""
    @function: 设置日志
"""
def setup_logger():

    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    
    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                                  datefmt="%Y-%m-%d %H:%M:%S")
    
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)

    log_folder = "log"
    os.makedirs(log_folder,exist_ok=True)

    current_date = datetime.now().strftime("%Y-%m-%d")
    log_file_name = f"log_{current_date}.txt"
    log_hander = logging.FileHandler(os.path.join(log_folder,log_file_name),mode='a',encoding='utf-8')

    log_hander.setLevel(logging.INFO)
    log_hander.setFormatter(formatter)


    logger.addHandler(console_handler)
    logger.addHandler(log_hander)

    return logger

"""
    @function: 初始化driver
"""
def init_driver():
    service = Service(r"C:\Users\86157\Desktop\PER_DEV\python\auto_deepseek\driver\chromedriver.exe")  # 直接指定驱动路径
    driver = webdriver.Chrome(service=service)
    driver.get("https://chat.deepseek.com/")
    return driver

"""
    @function: 初始化参数
"""
def init_parser():
    parser = argparse.ArgumentParser(description="一个浏览器自动化脚本，问答deepseek")
    parser.add_argument("--phone", required=True, type=str, help="手机号")
    parser.add_argument("--password", required=True, type=str, help="密码")

    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--question", type=str, help="问题")
    group.add_argument("--question_file", type=str, help="问题文件")
   
    parser.add_argument("--answer_file", type=str,help="回答文件")
    parser.add_argument("--deepseek", action="store_true", help="启用deepseek")
    parser.add_argument("--internet",action="store_true", help="联网搜索")

    parser.add_argument("--stable_index",type=float,default=1, help="稳健指数，指数越大，耗时越长，确保监听的数据越完整，建设数值为(0.5-2)")
    args = parser.parse_args()
    return args



if __name__ == "__main__":
    logger = setup_logger()
    args = init_parser()


"""
    @function: 登录
    @param: phone,password
    @return: boolean
"""
def auto_login(driver,phone,password,stable_index=1):
    wait = WebDriverWait(driver, timeout=10*stable_index)
    try:
        login_type = wait.until(
            EC.visibility_of_all_elements_located((By.CLASS_NAME, "ds-tab"))  
        )
        login_type[1].click()
        phone_input = wait.until(
            EC.visibility_of_element_located(
                (By.CSS_SELECTOR, "input[type='text'][placeholder='请输入手机号/邮箱地址']")
            )
        )
        phone_input.send_keys(phone)
        password_input = wait.until(
            EC.visibility_of_element_located(
                (By.CSS_SELECTOR, "input[type='password'][placeholder='请输入密码']")
            )
        )
        password_input.send_keys(password)
    
        confirm = wait.until(
            EC.visibility_of_element_located(
                (By.CSS_SELECTOR, ".ds-checkbox.ds-checkbox--none.ds-checkbox--bordered")
            )
        )
        confirm.click()

        login_button = wait.until(
            EC.visibility_of_element_located(
                (By.CSS_SELECTOR, ".ds-button.ds-button--primary.ds-sign-up-form__register-button")
            )
        )
        login_button.click()

        # if(wait.until(
        #     EC.visibility_of_element_located(
        #         (By.XPATH,"//div[text()='手机号码/邮箱地址或密码错误，请重新填写']"))
        #     )
        # ):
        #     raise Exception("手机号码/邮箱地址或密码错误，请重新填写")

        return True
    except TimeoutException:
        logger.error(f"元素未在指定时间内可见！")
        return False
    except Exception as e:
        logger.error(e)
        return False


"""
    @function: 获取思考内容
    @return: list
"""
def listener_think(driver, max_attempts=100,stable_count=5,stable_index=1):
    wait = WebDriverWait(driver, timeout=5*stable_index)
    attempts = 0
    while attempts < max_attempts*stable_index:
        try:
            content_div = wait.until(
                EC.visibility_of_element_located((By.XPATH,"//div[@class='e1675d8b']"))
            )
            break
        except TimeoutException:
            attempts += 1
            logger.info(f"第{attempts}次尝试获取思考内容")
            try:
                main_body = wait.until(
                    EC.visibility_of_element_located((By.XPATH,"//div[@class='ds-markdown ds-markdown--block']"))
                )
                if main_body:
                    p = main_body.find_element(By.XPATH,".//p")
                if(p.text == "服务器繁忙，请稍后再试。"):
                    logger.error(f"服务器繁忙，请稍后再试。")
                    return None
            except TimeoutException:
                pass
    
    previous_p_tags_len = 0
    check = 0
    stable_checks = 0
    while check < max_attempts*stable_index:
        check += 1
        try:
            wait.until(
                lambda d: len(content_div.find_elements(By.XPATH,".//p[normalize-space()]")) > previous_p_tags_len
            )
            logger.info(f"思考内容--段落更新中...")
        except TimeoutException:
            pass
        except StaleElementReferenceException:
            content_div = wait.until(
                EC.visibility_of_element_located((By.XPATH,"//div[@class='e1675d8b']"))
            )
                
        p_tags = content_div.find_elements(By.XPATH,".//p[normalize-space()]")
        current_p_tags_len = len(p_tags)

        if current_p_tags_len == previous_p_tags_len:
            stable_checks += 1
            if stable_checks >= stable_count*stable_index:
                logger.info(f"思考内容--段落更新完毕")
                break
        else:
            stable_checks = 0
            previous_p_tags_len = current_p_tags_len
                    
    think_list = []
    p_tags = content_div.find_elements(By.XPATH,".//p[normalize-space()]")
    for i in p_tags:
        think_list.append(i.text)
    return think_list


"""
    @function: 获取回答内容
    @return: list
"""         
def listener_answer(driver,max_attempts=100,stable_count=5,stable_index=1):
    wait = WebDriverWait(driver, timeout=5*stable_index)
    attempts = 0
    while attempts < max_attempts * stable_index:
        try:
            content_div = wait.until(
                EC.visibility_of_element_located((By.XPATH,"//div[@class='ds-markdown ds-markdown--block']"))
            )
            break
        except TimeoutException:
            attempts += 1
            logger.info(f"第{attempts}次尝试获取回答")
    
    previous_p_tags_len = 0
    check = 0
    stable_checks = 0
    while check < max_attempts * stable_index:
        check += 1
        try:
            wait.until(
                lambda d: len(content_div.find_elements(By.XPATH,".//p[normalize-space()]")) > previous_p_tags_len
            )
            logger.info(f"回答内容--段落更新中...")
        except TimeoutException:
            pass
        except StaleElementReferenceException:
            content_div = wait.until(
                EC.visibility_of_element_located((By.XPATH,"//div[@class='ds-markdown ds-markdown--block']"))
            )

        p_tags = content_div.find_elements(By.XPATH,".//p[normalize-space()]")
        current_p_tags_len = len(p_tags)
        if current_p_tags_len == previous_p_tags_len:
            stable_checks += 1
            if stable_checks >= stable_count * stable_index:
                logger.info(f"回答内容--段落更新完毕")
                break
        else:
            stable_checks = 0
            previous_p_tags_len = current_p_tags_len
    
    main_body_list = []
    p_tags = content_div.find_elements(By.XPATH,".//p[normalize-space()]")
    for i in p_tags:
        main_body_list.append(i.text)
    return main_body_list


"""
    @function: 启动深度思考
    @return: boolean
"""
def start_deepseek(driver,stable_index=1):
    try:
        wait = WebDriverWait(driver, timeout=5 * stable_index)
        deepseek_button = wait.until(
             EC.visibility_of_element_located((By.XPATH, "//span[text()='深度思考 (R1)']"))
        )
        deepseek_button.click()
        return True
    except Exception as e:
        logger.error(e)
        return False

"""
    @function: 启动联网搜索
    @return: boolean
"""
def start_internet_search(driver,stable_index=1):
    wait = WebDriverWait(driver, timeout=5 * stable_index)
    try:
        internet_search_button = wait.until(
                EC.visibility_of_element_located((By.XPATH, "//span[text()='联网搜索']"))
        )
        internet_search_button.click()
        return True
    except Exception as e:
        logger.error(e)
        return False

"""
    @function: 自动发送问题
    @return: think_result,answer_result
"""   
def auto_one_send(driver,question,start_deepseek=False,stable_index=1):
    wait = WebDriverWait(driver, timeout=10 * stable_index)
    try:    
        chat = wait.until(
            EC.visibility_of_element_located((By.ID,"chat-input"))
        )
        chat.send_keys(question)

        send_button = wait.until(
            EC.element_to_be_clickable((By.XPATH,"//div[@role='button' and @aria-disabled = 'false']"))
        )
        send_button.click()

        logger.info(f"问题--{question}")

        think_result = []
        if start_deepseek:
            think_result = listener_think(driver=driver,stable_index=stable_index)
        answer_result = listener_answer(driver=driver,stable_index=stable_index)

        return think_result,answer_result

    except ElementClickInterceptedException:
        logger.error("元素被遮挡")
    except TimeoutError:
        logger.error("元素未在指定时间内可见！")
    

"""
    @function: 上传文件
    @return: boolean
"""
def upload_file(file_path, time_limit=60,stable_index=1):
    file_name = os.path.basename(file_path)
    wait = WebDriverWait(driver, timeout=10 * stable_index)
    file_input = wait.until(
        EC.presence_of_element_located((By.CSS_SELECTOR,"input[type='file'][multiple]"))
    )
    file_input.send_keys(file_path)
    try: 
        file_name_element = wait.until(
            EC.visibility_of_element_located((By.XPATH,f"//div[text()='{file_name}']"))
        )
        # 定位文件状态元素的定位器
        file_status_locator = (By.XPATH,"./following-sibling::div[1]")
        

        WebDriverWait(driver,timeout=time_limit*stable_index).until(
            lambda d: file_name_element.find_element(*file_status_locator).text 
                        not in ['等待中...','上传中...','解析中...']
        )
        upload_file_status = file_name_element.find_element(*file_status_locator).text
        return True
        
    except TimeoutException:
        return False
        




phone = args.phone
password = args.password



question = args.question
question_file = args.question_file
stable_index = args.stable_index
if question:
    try:
        driver = init_driver()
        auto_login(driver=driver,phone=phone,password=password,stable_index=stable_index)

        if args.deepseek:
            start_deepseek(driver=driver)
        if args.internet:
            start_internet_search(driver=driver,stable_index=stable_index)

        think_result, answer_result = auto_one_send(driver=driver,question=question,start_deepseek=args.deepseek,stable_index=stable_index)

        think_result = "\n".join(think_result)
        print(f"思考过程如下: ")
        print(f"{think_result}")
        answer_result = "\n".join(answer_result)
        print(f"回答结果如下: ")
        print(f"{answer_result}")
    except Exception as e:
        logger.error(e)
    finally:
        driver.quit()
elif question_file:

    answer_folder_name = "answer"
    os.makedirs(answer_folder_name,exist_ok=True)

    answer_file_name = f"{os.path.splitext(os.path.basename(question_file))[0]}.txt"
    answer_file = os.path.join(answer_folder_name,answer_file_name)
    

    with open(question_file,"r",encoding="utf-8") as file:
        data = json.load(file)


    result_dict = {
        group_name: [item["question"] for item in questions]
        for group_name, questions in data.items()
    }
    
    with open(answer_file,"w",encoding="utf-8") as file:
        pass

    for group_name, questions in result_dict.items():
        try:
            with open(answer_file,"a", encoding="utf-8") as file:
                file.write(f"{group_name}组问题回答如下: \n")
            
            
            driver = init_driver()
            auto_login(driver=driver,phone=phone,password=password,stable_index=stable_index)
            if args.deepseek:
                start_deepseek(driver=driver,stable_index=stable_index)
            if args.internet:
                start_internet_search(driver=driver,stable_index=stable_index)
            
            logger.info(f"{group_name}组问题 开始思索回答: ")

            for idx,question in enumerate(questions,1):
                think_result, answer_result = auto_one_send(driver=driver,question=question,start_deepseek=args.deepseek,stable_index=stable_index)
                think_result = "\n".join(think_result)
                answer_result = "\n".join(answer_result)

                with open(answer_file,"a",encoding="utf-8") as file:
                    file.write(f"问题{idx}: {question}\n思考过程如下: \n{think_result}\n回答结果如下: \n{answer_result}\n")
            
            logger.info(f"{group_name}组问题 finished!")
        except Exception as e:
            logger.error(e)
        finally:
            driver.quit()


input("程序已完成，敲击“回车键”关闭程序")

